package com.yhc.common.util;

import cn.hutool.core.codec.Base64;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

public class SHAWithRSASigner {

    private static final Logger logger = LoggerFactory.getLogger(SHAWithRSASigner.class);

    private static final String ALGORITHM_RSA = "RSA";
    private static final String ALGORITHM_SHA256 = "SHA-256";
    private static final String ALGORITHM_SHA256WITHRSA = "SHA256withRSA";


    public static String sign(String privateKey, String plainText) {
        return sign(getPrivateKey(privateKey), plainText);
    }

    public static boolean verify(String publicKey, String plainText, String signed) {
        return verify(getPublicKey(publicKey), plainText, signed);
    }

    public static String sign(PrivateKey privateKey, String plainText) {
        try {
            MessageDigest messageDigest = MessageDigest.getInstance(ALGORITHM_SHA256);
            messageDigest.update(plainText.getBytes());
            byte[] outputDigest_sign = messageDigest.digest();

            Signature signature = Signature.getInstance(ALGORITHM_SHA256WITHRSA);
            signature.initSign(privateKey);
            signature.update(outputDigest_sign);
            byte[] signedData = signature.sign();
            return Base64.encode(signedData);
        } catch (NoSuchAlgorithmException | SignatureException e) {
            return null;
        } catch (InvalidKeyException e) {
            throw new IllegalArgumentException("[EX] illegal privateKey", e);
        }
    }

    public static boolean verify(PublicKey publicKey, String plainText, String signed) {
        try {
            MessageDigest messageDigest = MessageDigest.getInstance(ALGORITHM_SHA256);
            messageDigest.update(plainText.getBytes());
            byte[] outputDigest_sign = messageDigest.digest();

            Signature signature = Signature.getInstance(ALGORITHM_SHA256WITHRSA);
            signature.initVerify(publicKey);
            signature.update(outputDigest_sign);
            return signature.verify(Base64.decode(signed));
        } catch (NoSuchAlgorithmException | SignatureException e) {
            return false;
        } catch (InvalidKeyException e) {
            throw new IllegalArgumentException("[EX] illegal privateKey", e);
        }
    }

    private static PublicKey getPublicKey(String key) {
        try {
            byte[] byteKey = Base64.decode(key);
            X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(byteKey);
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
            return keyFactory.generatePublic(x509EncodedKeySpec);
        } catch (Exception e) {
            logger.error("get public key error", e);
        }
        return null;
    }

    private static PrivateKey getPrivateKey(String key) {
        try {
            byte[] byteKey = Base64.decode(key);
            PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(byteKey);
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
            return keyFactory.generatePrivate(pkcs8EncodedKeySpec);
        } catch (Exception e) {
            logger.error("get private key error", e);
        }
        return null;
    }

    public static Map<String, Object> getKey() throws Exception {
        KeyPairGenerator keygen = KeyPairGenerator.getInstance(ALGORITHM_RSA);
        SecureRandom secureRandom = new SecureRandom();
        keygen.initialize(2048, secureRandom);
        KeyPair keys = keygen.genKeyPair();
        PrivateKey privateKey = keys.getPrivate();
        PublicKey publicKey = keys.getPublic();

        logger.info("public key:\n{}", Base64.encode(publicKey.getEncoded()));
        logger.info("private key:\n{}", Base64.encode(privateKey.getEncoded()));

        Map<String, Object> map = new HashMap<>();
        map.put("privateKey", privateKey);
        map.put("publicKey", publicKey);
        map.put("privateKeyStr", Base64.encode(privateKey.getEncoded()));
        map.put("publicKeyStr", Base64.encode(publicKey.getEncoded()));
        return map;
    }

    public static void main(String[] args) throws Exception {
        Map<String, Object> map = getKey();
        PrivateKey privateKey = (PrivateKey) map.get("privateKey");
        PublicKey publicKey = (PublicKey) map.get("publicKey");
        String privateKeyStr = (String) map.get("privateKeyStr");
        String publicKeyStr = (String) map.get("publicKeyStr");

        String plainText = "test123";

        String signed = sign(privateKey, plainText);
        logger.info("验签1签名：{}", signed);
        boolean signFlag = verify(publicKey, plainText, signed);
        logger.info("验签1是否成功：{}", signFlag);

        signed = sign(privateKeyStr, plainText);
        logger.info("验签2签名：{}", signed);
        signFlag = verify(publicKeyStr, plainText, signed);
        logger.info("验签2是否成功：{}", signFlag);
    }


}
